/**
 * Copyright (C) 2024 AIDC-AI
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.alibaba.langengine.tinkerpop.vectorstore;

import com.alibaba.langengine.core.embeddings.Embeddings;
import com.alibaba.langengine.core.indexes.Document;
import com.alibaba.langengine.tinkerpop.vectorstore.service.*;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.Disabled;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;

import java.util.*;

import static org.junit.jupiter.api.Assertions.*;
import static org.mockito.Mockito.*;


@ExtendWith(MockitoExtension.class)
public class TinkerPopMockTest {

    @Mock
    private Embeddings mockEmbedding;

    @Mock
    private TinkerPopService mockService;

    @Mock
    private TinkerPopClient mockClient;

    private TinkerPop tinkerPop;
    private final String TEST_COLLECTION_ID = "test_collection";
    private final String TEST_SERVER_URL = "ws://test-server:8182";

    @BeforeEach
    public void setUp() {
        // 创建一个 TinkerPop 实例并设置 mock 服务
        tinkerPop = new TinkerPop(mockEmbedding, TEST_COLLECTION_ID);

        // 使用反射设置 mock 服务
        setPrivateField(tinkerPop, "_service", mockService);
        setPrivateField(tinkerPop, "_client", mockClient);
    }

    @AfterEach
    public void tearDown() {
        if (tinkerPop != null) {
            tinkerPop.close();
        }
    }

    @Test
    public void testConstructorWithNullCollectionId() {
        TinkerPop tinkerPopWithNullId = new TinkerPop(mockEmbedding, null);
        assertNotNull(tinkerPopWithNullId.getCollectionId());
        assertNotEquals("", tinkerPopWithNullId.getCollectionId());
        tinkerPopWithNullId.close();
    }

    @Test
    @Disabled("Requires a running TinkerPop server at test-server:8182")
    public void testConstructorWithCustomServerUrl() {
        TinkerPop customTinkerPop = new TinkerPop(TEST_SERVER_URL, mockEmbedding, TEST_COLLECTION_ID);
        assertNotNull(customTinkerPop);
        assertEquals(TEST_COLLECTION_ID, customTinkerPop.getCollectionId());
        assertEquals(mockEmbedding, customTinkerPop.getEmbedding());
        customTinkerPop.close();
    }

    @Test
    @Disabled("Requires a running TinkerPop server at test-server:8182")
    public void testConstructorWithTimeouts() {
        TinkerPop timeoutTinkerPop = new TinkerPop(
            TEST_SERVER_URL, mockEmbedding, TEST_COLLECTION_ID, 5000, 10000);
        assertNotNull(timeoutTinkerPop);
        assertEquals(TEST_COLLECTION_ID, timeoutTinkerPop.getCollectionId());
        timeoutTinkerPop.close();
    }

    @Test
    public void testAddDocumentsWithValidDocuments() {
        // 准备测试数据
        Document doc1 = new Document();
        doc1.setUniqueId("doc1");
        doc1.setPageContent("Test content 1");
        doc1.setEmbedding(Arrays.asList(0.1, 0.2, 0.3));
        doc1.setMetadata(Map.of("source", "test1.txt"));

        Document doc2 = new Document();
        doc2.setUniqueId("doc2");
        doc2.setPageContent("Test content 2");
        doc2.setEmbedding(Arrays.asList(0.4, 0.5, 0.6));
        doc2.setMetadata(Map.of("source", "test2.txt"));

        List<Document> documents = Arrays.asList(doc1, doc2);

        // Mock 服务行为
        doNothing().when(mockService).addDocuments(any(TinkerPopAddRequest.class));

        // 执行测试
        assertDoesNotThrow(() -> tinkerPop.addDocuments(documents));

        // 验证服务被调用
        verify(mockService).addDocuments(any(TinkerPopAddRequest.class));
    }

    @Test
    public void testAddDocumentsWithEmptyContent() {
        // 准备测试数据 - 包含空内容
        Document docWithEmptyContent = new Document();
        docWithEmptyContent.setUniqueId("empty-doc");
        docWithEmptyContent.setPageContent(""); // 空内容

        Document docWithNullContent = new Document();
        docWithNullContent.setUniqueId("null-doc");
        docWithNullContent.setPageContent(null); // null内容

        List<Document> documents = Arrays.asList(docWithEmptyContent, docWithNullContent);

        doNothing().when(mockService).addDocuments(any(TinkerPopAddRequest.class));

        assertDoesNotThrow(() -> tinkerPop.addDocuments(documents));
        verify(mockService).addDocuments(any(TinkerPopAddRequest.class));
    }

    @Test
    public void testAddDocumentsWithAutoGeneratedIds() {
        Document docWithoutId = new Document();
        docWithoutId.setPageContent("Content without ID");

        List<Document> documents = Arrays.asList(docWithoutId);

        doNothing().when(mockService).addDocuments(any(TinkerPopAddRequest.class));

        assertDoesNotThrow(() -> tinkerPop.addDocuments(documents));

        // 验证ID被自动生成
        assertNotNull(docWithoutId.getUniqueId());
        assertFalse(docWithoutId.getUniqueId().isEmpty());

        verify(mockService).addDocuments(any(TinkerPopAddRequest.class));
    }

    @Test
    public void testAddDocumentsWithEmbeddingGeneration() {
        Document docWithoutEmbedding = new Document();
        docWithoutEmbedding.setUniqueId("doc-no-embedding");
        docWithoutEmbedding.setPageContent("Content without embedding");

        List<Document> documents = Arrays.asList(docWithoutEmbedding);

        // Mock embedding 服务生成嵌入向量
        Document embeddedDoc = new Document();
        embeddedDoc.setEmbedding(Arrays.asList(0.1, 0.2, 0.3));
        when(mockEmbedding.embedTexts(anyList())).thenReturn(Arrays.asList(embeddedDoc));

        doNothing().when(mockService).addDocuments(any(TinkerPopAddRequest.class));

        assertDoesNotThrow(() -> tinkerPop.addDocuments(documents));

        verify(mockEmbedding).embedTexts(anyList());
        verify(mockService).addDocuments(any(TinkerPopAddRequest.class));
    }

    @Test
    public void testAddTexts() {
        List<String> texts = Arrays.asList("Text 1", "Text 2", "Text 3");
        List<Map<String, Object>> metadatas = Arrays.asList(
            Map.of("source", "file1.txt"),
            Map.of("source", "file2.txt"),
            Map.of("source", "file3.txt")
        );
        List<String> ids = Arrays.asList("id1", "id2", "id3");

        // Mock embedding 服务
        List<Document> embeddedDocs = Arrays.asList(
            createMockDocument("id1", Arrays.asList(0.1, 0.2)),
            createMockDocument("id2", Arrays.asList(0.3, 0.4)),
            createMockDocument("id3", Arrays.asList(0.5, 0.6))
        );
        when(mockEmbedding.embedTexts(texts)).thenReturn(embeddedDocs);

        doNothing().when(mockService).addDocuments(any(TinkerPopAddRequest.class));

        List<String> resultIds = tinkerPop.addTexts(texts, metadatas, ids);

        assertEquals(ids, resultIds);
        verify(mockEmbedding).embedTexts(texts);
        verify(mockService).addDocuments(any(TinkerPopAddRequest.class));
    }

    @Test
    public void testAddTextsWithNullIds() {
        List<String> texts = Arrays.asList("Text 1", "Text 2");
        List<Map<String, Object>> metadatas = Arrays.asList(
            Map.of("source", "file1.txt"),
            Map.of("source", "file2.txt")
        );

        // Mock embedding 服务
        List<Document> embeddedDocs = Arrays.asList(
            createMockDocument("auto-id-1", Arrays.asList(0.1, 0.2)),
            createMockDocument("auto-id-2", Arrays.asList(0.3, 0.4))
        );
        when(mockEmbedding.embedTexts(texts)).thenReturn(embeddedDocs);

        doNothing().when(mockService).addDocuments(any(TinkerPopAddRequest.class));

        List<String> resultIds = tinkerPop.addTexts(texts, metadatas, null);

        assertEquals(2, resultIds.size());
        // 验证ID被自动生成
        resultIds.forEach(id -> {
            assertNotNull(id);
            assertFalse(id.isEmpty());
        });

        verify(mockEmbedding).embedTexts(texts);
        verify(mockService).addDocuments(any(TinkerPopAddRequest.class));
    }

    @Test
    public void testSimilaritySearch() {
        String query = "test query";
        int k = 5;

        // Mock embedding 服务为查询生成嵌入向量
        Document queryEmbedDoc = createMockDocument("query", Arrays.asList(0.1, 0.2, 0.3));
        when(mockEmbedding.embedTexts(Arrays.asList(query))).thenReturn(Arrays.asList(queryEmbedDoc));

        // Mock 服务查询响应
        TinkerPopQueryResponse mockResponse = new TinkerPopQueryResponse();
        mockResponse.setIds(Arrays.asList("doc1", "doc2"));
        mockResponse.setTexts(Arrays.asList("result text 1", "result text 2"));
        mockResponse.setDistances(Arrays.asList(0.1, 0.2));
        mockResponse.setMetadatas(Arrays.asList(
            Map.of("source", "file1.txt"),
            Map.of("source", "file2.txt")
        ));

        when(mockService.queryDocuments(any(TinkerPopQueryRequest.class))).thenReturn(mockResponse);

        List<Document> results = tinkerPop.similaritySearch(query, k, null, null);

        assertEquals(2, results.size());
        assertEquals("doc1", results.get(0).getUniqueId());
        assertEquals("doc2", results.get(1).getUniqueId());
        assertEquals("result text 1", results.get(0).getPageContent());
        assertEquals("result text 2", results.get(1).getPageContent());

        verify(mockEmbedding).embedTexts(Arrays.asList(query));
        verify(mockService).queryDocuments(any(TinkerPopQueryRequest.class));
    }

    @Test
    public void testSimilaritySearchWithMaxDistance() {
        String query = "test query";
        int k = 5;
        Double maxDistance = 0.15;

        // Mock embedding 服务
        Document queryEmbedDoc = createMockDocument("query", Arrays.asList(0.1, 0.2, 0.3));
        when(mockEmbedding.embedTexts(Arrays.asList(query))).thenReturn(Arrays.asList(queryEmbedDoc));

        // Mock 服务查询响应 - 包含距离超过阈值的结果
        TinkerPopQueryResponse mockResponse = new TinkerPopQueryResponse();
        mockResponse.setIds(Arrays.asList("doc1", "doc2", "doc3"));
        mockResponse.setTexts(Arrays.asList("result 1", "result 2", "result 3"));
        mockResponse.setDistances(Arrays.asList(0.1, 0.2, 0.3)); // doc2 和 doc3 超过阈值
        mockResponse.setMetadatas(Arrays.asList(
            Map.of("source", "file1.txt"),
            Map.of("source", "file2.txt"),
            Map.of("source", "file3.txt")
        ));

        when(mockService.queryDocuments(any(TinkerPopQueryRequest.class))).thenReturn(mockResponse);

        List<Document> results = tinkerPop.similaritySearch(query, k, maxDistance, null);

        // 只有 doc1 的距离 <= 0.15
        assertEquals(1, results.size());
        assertEquals("doc1", results.get(0).getUniqueId());
        assertEquals(0.1, results.get(0).getScore());

        verify(mockEmbedding).embedTexts(Arrays.asList(query));
        verify(mockService).queryDocuments(any(TinkerPopQueryRequest.class));
    }

    @Test
    public void testSimilaritySearchWithNullEmbedding() {
        tinkerPop.setEmbedding(null); // 设置为 null
        String query = "test query";
        int k = 5;

        TinkerPopQueryResponse mockResponse = new TinkerPopQueryResponse();
        mockResponse.setIds(Arrays.asList("doc1"));
        mockResponse.setTexts(Arrays.asList("result text"));
        mockResponse.setDistances(Arrays.asList(0.1));
        mockResponse.setMetadatas(Arrays.asList(Map.of("source", "file.txt")));

        when(mockService.queryDocuments(any(TinkerPopQueryRequest.class))).thenReturn(mockResponse);

        List<Document> results = tinkerPop.similaritySearch(query, k, null, null);

        assertEquals(1, results.size());
        verify(mockService).queryDocuments(any(TinkerPopQueryRequest.class));
        // 验证没有调用 embedding 服务
        verifyNoInteractions(mockEmbedding);
    }

    @Test
    public void testSimilaritySearchWithHtmlFiltering() {
        String query = "test query";

        Document queryEmbedDoc = createMockDocument("query", Arrays.asList(0.1, 0.2, 0.3));
        when(mockEmbedding.embedTexts(Arrays.asList(query))).thenReturn(Arrays.asList(queryEmbedDoc));

        // Mock 响应包含 HTML 标签的文本
        TinkerPopQueryResponse mockResponse = new TinkerPopQueryResponse();
        mockResponse.setIds(Arrays.asList("doc1"));
        mockResponse.setTexts(Arrays.asList("<p>This is <b>bold</b> text &amp; escaped</p>"));
        mockResponse.setDistances(Arrays.asList(0.1));
        mockResponse.setMetadatas(Arrays.asList(Map.of("source", "file.txt")));

        when(mockService.queryDocuments(any(TinkerPopQueryRequest.class))).thenReturn(mockResponse);

        List<Document> results = tinkerPop.similaritySearch(query, 5, null, null);

        assertEquals(1, results.size());
        // 验证 HTML 标签被过滤，HTML 实体被解码
        assertEquals("This is bold text & escaped", results.get(0).getPageContent());
    }

    @Test
    public void testAddDocumentsThrowsException() {
        Document doc = new Document();
        doc.setUniqueId("doc1");
        doc.setPageContent("content");
        doc.setMetadata(Map.of("key", "value"));

        doThrow(new RuntimeException("Service error")).when(mockService).addDocuments(any(TinkerPopAddRequest.class));

        RuntimeException exception = assertThrows(RuntimeException.class, () -> {
            tinkerPop.addDocuments(Arrays.asList(doc));
        });

        assertTrue(exception.getMessage().contains("Service error"));
    }

    @Test
    public void testSimilaritySearchThrowsException() {
        when(mockService.queryDocuments(any(TinkerPopQueryRequest.class)))
            .thenThrow(new RuntimeException("Query error"));

        RuntimeException exception = assertThrows(RuntimeException.class, () -> {
            tinkerPop.similaritySearch("query", 5, null, null);
        });

        assertTrue(exception.getMessage().contains("Query error"));
    }

    @Test
    public void testClose() {
        doNothing().when(mockService).close();

        assertDoesNotThrow(() -> tinkerPop.close());
        verify(mockService).close();
    }

    @Test
    public void testFilterMethod() {
        // 使用反射测试私有的 filter 方法
        String htmlText = "<p>Test <b>bold</b> text &lt;script&gt; &amp; more</p>";
        String filtered = invokeFilterMethod(htmlText);

        assertEquals("Test bold text <script> & more", filtered);
    }

    @Test
    public void testFilterWithNullInput() {
        String filtered = invokeFilterMethod(null);
        assertEquals("", filtered);
    }

    // 辅助方法
    private Document createMockDocument(String id, List<Double> embedding) {
        Document doc = new Document();
        doc.setUniqueId(id);
        doc.setEmbedding(embedding);
        return doc;
    }

    private void setPrivateField(Object object, String fieldName, Object value) {
        try {
            java.lang.reflect.Field field = object.getClass().getDeclaredField(fieldName);
            field.setAccessible(true);
            field.set(object, value);
        } catch (Exception e) {
            fail("Failed to set private field: " + fieldName);
        }
    }

    private String invokeFilterMethod(String input) {
        try {
            java.lang.reflect.Method method = TinkerPop.class.getDeclaredMethod("filter", String.class);
            method.setAccessible(true);
            return (String) method.invoke(tinkerPop, input);
        } catch (Exception e) {
            fail("Failed to invoke filter method");
            return null;
        }
    }
}